#include <memory>
#include <iostream>
#include <string>
#include <vector>
#include <cmath>

#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/point_cloud2.hpp>

// PCL includes
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_types.h>
#include <pcl/common/transforms.h>
#include <pcl/filters/extract_indices.h>
#include <pcl/filters/passthrough.h>
#include <pcl/features/normal_3d.h>
#include <pcl/search/kdtree.h>
#include <pcl/filters/voxel_grid.h>

using std::placeholders::_1;

class FloorFilterNode : public rclcpp::Node {
public:
  using PointT = pcl::PointXYZI;

  FloorFilterNode() : Node("floor_filter_node") {
    // --- 参数声明 ---
    // 粗略的高度过滤范围 (只处理 Z 在这个范围内的点，减少计算量)
    this->declare_parameter<double>("min_z", -2.0);
    this->declare_parameter<double>("max_z", 2.0);
    
    // 法向量计算的搜索半径 (米)
    // 越大越平滑，但对细小台阶不敏感；越小细节越多，但容易受噪声影响
    this->declare_parameter<double>("normal_search_radius", 0.5);
    
    // 允许的最大坡度角 (度)
    // 0度 = 绝对水平，30度 = 可以包含斜坡，90度 = 包含垂直墙面
    this->declare_parameter<double>("max_slope_angle", 25.0);

    // 降采样 (米) - 既然"铺不满"，稍微降采样可以让点分布更均匀
    this->declare_parameter<double>("leaf_size", 0.1);

    min_z_ = this->get_parameter("min_z").as_double();
    max_z_ = this->get_parameter("max_z").as_double();
    normal_search_radius_ = this->get_parameter("normal_search_radius").as_double();
    max_slope_angle_ = this->get_parameter("max_slope_angle").as_double();
    leaf_size_ = this->get_parameter("leaf_size").as_double();

    RCLCPP_INFO(this->get_logger(), "Method: Normal Filtering. Slope < %.1f deg", max_slope_angle_);

    // QoS: 必须是 Transient Local (Latched)
    rclcpp::QoS map_qos_profile(1);
    map_qos_profile.transient_local();
    map_qos_profile.reliable();

    sub_points_ = this->create_subscription<sensor_msgs::msg::PointCloud2>(
      "input_cloud", 
      map_qos_profile, 
      std::bind(&FloorFilterNode::cloud_callback, this, _1));

    pub_ground_ = this->create_publisher<sensor_msgs::msg::PointCloud2>(
      "output_ground", 
      map_qos_profile);
  }

private:
  void cloud_callback(const sensor_msgs::msg::PointCloud2::SharedPtr msg) {
    RCLCPP_INFO(this->get_logger(), "Processing map...");

    pcl::PointCloud<PointT>::Ptr cloud(new pcl::PointCloud<PointT>());
    pcl::fromROSMsg(*msg, *cloud);

    if (cloud->empty()) return;

    // 1. 降采样 (Uniform Sampling)
    // 这步很重要！它可以让地面点云密度均匀，填补一些细小的空洞，同时加速法向量计算
    pcl::PointCloud<PointT>::Ptr cloud_ds(new pcl::PointCloud<PointT>());
    pcl::VoxelGrid<PointT> sor;
    sor.setInputCloud(cloud);
    sor.setLeafSize(leaf_size_, leaf_size_, leaf_size_);
    sor.filter(*cloud_ds);

    // 2. 粗略高度过滤 (PassThrough)
    // 先把明显的天花板和地底噪点去掉
    pcl::PointCloud<PointT>::Ptr cloud_z_filtered(new pcl::PointCloud<PointT>());
    pcl::PassThrough<PointT> pass;
    pass.setInputCloud(cloud_ds);
    pass.setFilterFieldName("z");
    pass.setFilterLimits(min_z_, max_z_);
    pass.filter(*cloud_z_filtered);

    if (cloud_z_filtered->empty()) return;

    // 3. 计算法向量 (Normal Estimation)
    pcl::NormalEstimation<PointT, pcl::Normal> ne;
    ne.setInputCloud(cloud_z_filtered);
    pcl::search::KdTree<PointT>::Ptr tree(new pcl::search::KdTree<PointT>());
    ne.setSearchMethod(tree);
    pcl::PointCloud<pcl::Normal>::Ptr cloud_normals(new pcl::PointCloud<pcl::Normal>);
    ne.setRadiusSearch(normal_search_radius_); // 关键参数
    ne.compute(*cloud_normals);

    // 4. 根据角度筛选地面
    pcl::PointCloud<PointT>::Ptr ground_cloud(new pcl::PointCloud<PointT>());
    
    // Z轴方向 (0,0,1)
    Eigen::Vector3f axis_z(0.0, 0.0, 1.0);
    double cos_thresh = std::cos(max_slope_angle_ * M_PI / 180.0);

    for (size_t i = 0; i < cloud_normals->points.size(); ++i) {
        Eigen::Vector3f normal(
            cloud_normals->points[i].normal_x,
            cloud_normals->points[i].normal_y,
            cloud_normals->points[i].normal_z
        );

        // 计算点云法向量与Z轴的夹角余弦值
        // 使用 abs 是因为法向量可能指向 Z+ 也可能指向 Z-，只要垂直于地面即可
        double dot = std::abs(normal.dot(axis_z));

        // 如果夹角小于阈值 (即 dot 接近 1)，则认为是地面
        if (dot >= cos_thresh) {
            ground_cloud->points.push_back(cloud_z_filtered->points[i]);
        }
    }

    ground_cloud->width = ground_cloud->points.size();
    ground_cloud->height = 1;
    ground_cloud->is_dense = true;
    ground_cloud->header = cloud_ds->header;

    RCLCPP_INFO(this->get_logger(), "Ground extracted. Points: %ld", ground_cloud->size());
    
    sensor_msgs::msg::PointCloud2 output_msg;
    pcl::toROSMsg(*ground_cloud, output_msg);
    output_msg.header.frame_id = msg->header.frame_id;
    output_msg.header.stamp = msg->header.stamp;
    
    pub_ground_->publish(output_msg);
  }

  // Params
  double min_z_, max_z_;
  double normal_search_radius_;
  double max_slope_angle_;
  double leaf_size_;

  rclcpp::Subscription<sensor_msgs::msg::PointCloud2>::SharedPtr sub_points_;
  rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr pub_ground_;
};

int main(int argc, char** argv) {
  rclcpp::init(argc, argv);
  rclcpp::spin(std::make_shared<FloorFilterNode>());
  rclcpp::shutdown();
  return 0;
}